#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA
#

from DynamicSolverBase import DynamicSolverBase

import time
from scipy import asarray
from SecondOrderStatic1D import SecondOrderStatic1D
from utils.ProjOperators import ProjOnShapeFunctions


class CDRSolver(DynamicSolverBase):

    """
    Convection diffusion reaction solver.
    Finite element P1, P2, splines.
    Not stabilize.
    """

    __slots__ = ['m_diags', 'm_mass_bd', 'm_stiff_bd', 'm_sos',
                 'm_pival']

    def __init__(self, nsegs, xleft, xright):
        """
      m_pival ~ initial condition (projected on fem space)
    """

        (self.m_diags, self.m_pival) = (None, None)
        (self.m_mass_bd, self.m_stiff_bd) = (None, None)

        DynamicSolverBase.__init__(self)
        self.m_sos = SecondOrderStatic1D(nsegs, xleft, xright)

    def __bd_operator(self, dxl, dxr):
        """ boudary conditions """

        basis = self.m_sos.get_sfunc()

        diffu_mat = basis.get_grad_grad()
        conv_mat = basis.get_grad_disp()
        mass_mat = basis.get_mass_mat()
        mass_left_bd = mass_mat[1, 0] * dxl
        mass_right_bd = mass_mat[0, -1] * dxr
        mass_at_bd = asarray([mass_left_bd, mass_right_bd])

        (d, c, r) = self.m_sos.get_pdecos()

        left_kmat = -d * diffu_mat[1, 0] / dxl + c * conv_mat[1, 0] + r \
            * mass_left_bd
        right_kmat = -d * diffu_mat[0, -1] / dxr + c * conv_mat[0, -1] \
            + r * mass_right_bd
        stiff_at_bd = asarray([left_kmat, right_kmat])

        if 3 == basis.dof():
            mass_left_bd = mass_mat[2, 0] * dxl
            mass_right_bd = mass_mat[1, -1] * dxr
            mass_at_bd = asarray([mass_at_bd[0], mass_right_bd,
                                 mass_left_bd, mass_at_bd[-1]],
                                 'float64')

            left_kmat = -d * diffu_mat[2, 0] / dxl + c * conv_mat[2, 0] \
                + r * mass_left_bd
            right_kmat = -d * diffu_mat[1, -1] / dxr + c * conv_mat[1,
                    -1] + r * mass_right_bd
            stiff_at_bd = asarray([stiff_at_bd[0], right_kmat,
                                  left_kmat, stiff_at_bd[-1]], 'float64'
                                  )

        mass_at_bd = mass_at_bd.reshape((basis.dof() - 1, 2))
        stiff_at_bd = stiff_at_bd.reshape((basis.dof() - 1, 2))
        return (mass_at_bd, stiff_at_bd)

    def _time_step(self, mass_bd, stiff_bd, diags):
        """ set up time stepper, time step and return final result. """

        grid = self.m_sos.get_mesh()
        dlx = grid.getDx()

        (xleft, xright) = grid.getRange()
        (bd_mass, bd_kmat) = self.__bd_operator(dlx[0], dlx[-1])

        # Initialize Time stepper.

        tstepper = self.get_tstepper()
        tstepper.set_boperators(bd_mass, bd_kmat)

        tstepper.set_xlxr(xleft, xright)
        tstepper.set_ival(self.m_pival[1:-1])
        tstepper.set_bcond(self.get_lbc(), self.get_rbc())

        tstepper.set_moperators(mass_bd, stiff_bd, diags)

        dof = self.m_sos.get_sfunc().dof()
        if 2 == dof:
            tstepper.set_mat_attrib(3)
        if 3 == dof:
            tstepper.set_mat_attrib(5)

        return tstepper.step()

    def get_mesh(self):
        """ grid use to assemble matrices. """

        return self.m_sos.get_mesh()


    def get_sfunc(self):
        """ shape functions defining the finite element space. """

        return self.m_sos.get_sfunc()

    def get_pival(self):
        """ get temporal solution."""

        return self.m_pival[1:-1]

    def set_diffus_cos(self, diff, conv, reac, d_is_real, c_is_real,
            r_is_real):
        """ set diffusion coefficients. """

        self.m_sos.set_cos(
            diff,
            conv,
            reac,
            d_is_real,
            c_is_real,
            r_is_real,
            )

    def set_sfunc_and_quads(self, shfunts, quad):
        """ finite element space and quad formulas. """

        self.m_sos.set_sfunc_and_quads(shfunts, quad)

    def set_up_problem(self):
        """ Only Dirichlet bc are implemented."""

        start = time.time()
        self.m_pival = ProjOnShapeFunctions(self.m_sos.get_sfunc(),
                self.m_sos.get_quads(), self.m_sos.get_mesh(),
                self.get_ivalf())

        # *************************************************

        dof = self.m_sos.get_sfunc().dof()
        st_diag = self.m_sos.get_stiff_dformat()
        ma_diag = self.m_sos.get_mass_dformat()

        if 2 == dof:
            stiff_bd = asarray([list(st_diag[1:-1, 0].flatten()),
                               list(st_diag[1:-1, 1].flatten()),
                               list(st_diag[1:-1, 2].flatten())])
            mass_bd = asarray([list(ma_diag[1:-1, 0].flatten()),
                              list(ma_diag[1:-1, 1].flatten()),
                              list(ma_diag[1:-1, 2].flatten())])
            diags = asarray([0, 1, -1])
        if 3 == dof:
            stiff_bd = asarray([list(st_diag[1:-1, 0].flatten()),
                               list(st_diag[1:-1, 1].flatten()),
                               list(st_diag[1:-1, 2].flatten()),
                               list(st_diag[1:-1, 3].flatten()),
                               list(st_diag[1:-1, 4].flatten())])
            mass_bd = asarray([list(ma_diag[1:-1, 0].flatten()),
                              list(ma_diag[1:-1, 1].flatten()),
                              list(ma_diag[1:-1, 2].flatten()),
                              list(ma_diag[1:-1, 3].flatten()),
                              list(ma_diag[1:-1, 4].flatten())])
            diags = asarray([0, 1, -1, 2, -2])

        if self.do_log():
            self.log_info('')
            self.log_info('Assembling system takes ' + str(time.time()
                          - start) + ' secs')

        (self.m_mass_bd, self.m_stiff_bd, self.m_diags) = (mass_bd,
                stiff_bd, diags)

    def solve(self):
        """ solve pde """

        if None == self.m_pival:
            self.set_up_problem()

        return self._time_step(self.m_mass_bd, self.m_stiff_bd,
                               self.m_diags)


